"""
来自论文[1] Hanak Chad, Crain TImothy, Bishop Robert. Crater identification algorithm for the lost in low lunar orbit scenario[C]. 33rd Annual AAS Guidance and Control. Breckenridge, CO, US: American Astronautical Society, 2010

仅使用论文提供的不变量作为描述子，不使用该论文的投票法作为识别方法。确定搜索阈值的方法是：预先对当前数据库作一次蒙特卡洛采样，求得所有不变量的标准差，以3倍标准差作为搜索阈值，搜索阈值的确定应当在TriadBuilder中进行。
"""

import numpy as np
import torch
from typing import Iterator
from utils.ellipse import ellipse_args, center_uncertainty, axis_uncertainty
from .base import TriadPyramidVote
from itertools import combinations
from scipy.optimize import linear_sum_assignment


class HanakPyramidVote(TriadPyramidVote):

    def search_threshold(self, C, U, L, Cos, x, y, r1, r2):
        """计算搜索阈值"""
        u_x, u_y = center_uncertainty(C, U)
        u_a, u_b = axis_uncertainty(C, U, r1, r2)
        u_r = np.sqrt(u_a**2 + u_b**2) * 0.5
        delta_x = np.array((x[1] - x[2], x[0] - x[2], x[0] - x[1]))
        delta_y = np.array((y[1] - y[2], y[0] - y[2], y[0] - y[1]))
        delta_u_x = (u_x * u_x).sum() - u_x * u_x
        delta_u_y = (u_y * u_y).sum() - u_y * u_y
        u_l = 1 / L * np.sqrt(delta_x**2 * delta_u_x + delta_y**2 * delta_u_y)
        u_123 = 2 / L[0] * np.sqrt(u_r**2 + ((r1 + r2) / 2 * u_l[0] / L[0]) ** 2)
        pi_L = np.prod(L, axis=0)
        u_4 = (
            L[[0, 1]] ** 2
            / pi_L
            * np.sqrt(
                (
                    (Cos[2] * u_l[[1, 0]]) ** 2
                    + (Cos[[1, 0]] * u_l[2]) ** 2
                    + (u_l[[0, 1]]) ** 2
                )
            )
        )
        return np.pad((*u_123, *u_4), ((0, 1), (0, 0)), constant_values=0)

    def triad_descriptor(
        self, i, j, k, Q1, Q2, Q3, U1, U2, U3, *args, factor=0.3, **kwargs
    ):
        """从陨石坑的参数中计算描述子，C1和C2可以是含有不确定度的椭圆参数"""
        N = Q3.shape[0]
        if len(Q1.shape) == 2:
            Q1 = Q1[None].repeat(N, axis=0)
            U1 = U1[None].repeat(N, axis=0)
        if len(Q2.shape) == 2:
            Q2 = Q2[None].repeat(N, axis=0)
            U2 = U2[None].repeat(N, axis=0)
        assert Q1.shape[0] == Q2.shape[0] == Q3.shape[0]
        assert U1.shape[0] == U2.shape[0] == U3.shape[0]
        x1, y1, a1, b1, _ = ellipse_args(Q1)
        x2, y2, a2, b2, _ = ellipse_args(Q2)
        x3, y3, a3, b3, _ = ellipse_args(Q3)
        d1 = a1 + b1
        d2 = a2 + b2
        d3 = a3 + b3
        ds = np.array((d1, d2, d3))
        # 求解椭圆中心构成的三角形中两个较大的内角
        # 三角形的三个顶点坐标
        A = np.array((x1, y1, np.ones_like(x1)))
        B = np.array((x2, y2, np.ones_like(x1)))
        C = np.array((x3, y3, np.ones_like(x1)))
        ABC = np.array((A, B, C)).T
        # 三角形的三条边长
        a = np.linalg.norm(C - B, axis=0)
        b = np.linalg.norm(A - C, axis=0)
        c = np.linalg.norm(A - B, axis=0)
        l = np.max((a, b, c), axis=0)
        # 三角形的三个内角
        cosA = (b**2 + c**2 - a**2) / (2 * b * c)
        cosB = (a**2 + c**2 - b**2) / (2 * a * c)
        cosC = (a**2 + b**2 - c**2) / (2 * a * b)
        cos = np.array((cosA, cosB, cosC))
        # 按角度从大到小排序，即余弦值从小到大排序
        index = np.argsort(cos, axis=0)
        # 排除共线的情况，即三个点构成的行列式为0
        sequence = np.linalg.det(ABC)
        valid_ind = ~np.isclose(sequence, 0)
        # 排除椭圆大小重叠、不同椭圆交叉的情况
        valid_ind &= (c > (d1 + d2) / 2) & (b > (d1 + d3) / 2) & (a > (d2 + d3) / 2)
        # 三角形的两个较大内角，对应较小的两条边
        idx = np.arange(N)
        # 按原始顺序排序输出：三个对应椭圆半径与最大边长之比，最大内角、第二大内角的余弦值，以及按角度最大、第二大和最小顺序的旋转方向。
        ds = ds[index, idx]
        cos = cos[index, idx]
        I = np.concatenate((ds / l, cos[[0, 1]], np.sign([sequence])), axis=0).astype(
            np.float32
        )
        # 计算搜索阈值
        Q = np.array((Q1, Q2, Q3))[index, idx]
        U = np.array((U1, U2, U3))[index, idx]
        x = np.array((x1, x2, x3))[index, idx]
        y = np.array((y1, y2, y3))[index, idx]
        r1 = np.array((a1, a2, a3))[index, idx]
        r2 = np.array((b1, b2, b3))[index, idx]
        L = np.array((a, b, c))[index, idx]
        ## 利用不确定度计算搜索阈值
        th = self.search_threshold(Q, U, L, cos, x, y, r1, r2)
        ## 对于较小的不确定度，将其替换为浮动阈值
        low_ind = th < 1e-3
        th[low_ind] = np.abs(I[low_ind]) * factor
        ## 保证最后一行为0
        th[-1] = 0
        # 计算当前顺序
        ijk = np.vstack((i, j, k))
        # 返回的index必须是恢复原始顺序的
        return ijk, I, th, np.argsort(index, axis=0), valid_ind

    def identify(
        self, params, uncertainty, factor, *args, max_try=5, **kwargs
    ) -> Iterator:
        """
        输入应该越接近于真实的检测输出越好。因此输入应当是一堆椭圆的拟合参数，输出应当是匹配确定的结果，即椭圆ID的列表。
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        # 将陨石坑按直径排序
        params = np.array(params)
        elli = ellipse_args(params)
        diameters = elli[2] + elli[3]
        idx = np.argsort(-diameters)
        if len(idx) < 4:
            return None
        # 不要构建全部三元组，极其费时间
        # 从最大的直径开始，选取陨石坑构建金字塔
        fail_num = 0
        for i, j in combinations(range(idx.shape[0]), 2):
            k, l = np.array(list(combinations(np.delete(idx, [i, j]), 2))).T
            i = idx[i].repeat(k.shape[0])
            j = idx[j].repeat(k.shape[0])
            # for k, l in combinations(not_used[2:], 2):
            # 批量算出金字塔的三元组
            ijk, I123, th, index, valid = self.triad_descriptor(
                np.concatenate((i, i, j, i)),
                np.concatenate((j, j, k, k)),
                np.concatenate((k, l, l, l)),
                params[np.concatenate((i, i, j, i))],
                params[np.concatenate((j, j, k, k))],
                params[np.concatenate((k, l, l, l))],
                uncertainty[np.concatenate((i, i, j, i))],
                uncertainty[np.concatenate((j, j, k, k))],
                uncertainty[np.concatenate((k, l, l, l))],
                factor,
            )
            I123 = torch.tensor(I123, device=self.device, dtype=torch.float32)
            th = torch.tensor(th, device=self.device, dtype=torch.float32)
            # 变成四个三元组，即金字塔的结构
            ijk = ijk.reshape(3, 4, -1)
            I123 = I123.view(I123.shape[0], 4, -1)
            th = th.view(th.shape[0], 4, -1)
            index = index.reshape(3, 4, -1)
            valid = valid.reshape(4, -1).all(axis=0)
            if not valid.any():
                fail_num += 1
                continue
            # 由于上一步是批量计算的，这一步再批量显存要爆了，所以就一行一行的来吧
            ijkl, IJKL = self.vote_record(
                ijk[..., valid],
                I123[..., valid],
                th[..., valid],
                index[..., valid],
                n=i.shape[0],
            )
            if len(ijkl) > 4:
                # 成功一次后，直接用当前四点扩展
                # success, result = self.extend_identify(
                #     params, uncertainty, ijkl, IJKL, *args, factor=factor, **kwargs
                # )
                # if success:
                return ijkl, IJKL
            fail_num += 1
            if fail_num > max_try:
                break
        return None

    def vote_record(self, ijk, I123, th, index, n, vote_th=6):
        """
        投票法确定真实的匹配结果
        Arguments:
            params (np.ndarray) : ellipse parameters for (x^2, xy, y^2, x, y, 1)
        Returns:
            list : matched ellipse ID
        """
        votes = torch.zeros(
            (self.catalog.shape[0], self.catalog.shape[0]),
            dtype=np.int32,
            device=self.device,
        )
        # idx = np.arange(4, dtype=np.int32)
        for it in range(ijk.shape[2]):
            ABC = self.triad_search(I123[..., it], index[..., it], th[..., it])
            if ABC is None:
                continue
            pyra = ijk[..., it]
            assert np.unique(pyra[2, 1:]).shape[0] == 1
            ijkl = np.array((*pyra[:, 0], pyra[2, 1]))
            for IJKL in self.build_pyramids(ABC):
                # 完成一次全图扩展
                votes[ijkl, IJKL] += 1
        # 选取最大值
        votes = votes.cpu().numpy()
        ## 用匈牙利算法完成识别分配
        ijkl, IJKL = linear_sum_assignment(-votes)
        # 最优分配
        ind = votes[ijkl, IJKL] > vote_th
        return ijkl[ind], IJKL[ind]
